from dotenv import load_dotenv
import os

from langchain_community.vectorstores import Chroma

from model_utils import getEmd
from langchain_community.document_loaders import TextLoader
from zai import ZhipuAiClient
load_dotenv()

def get_web_search(query):
    client = ZhipuAiClient(api_key=os.getenv("ZHIPU_KEY"))

    response = client.web_search.web_search(
        search_engine="search_pro",
        search_query=query,
        count=1,  # 返回结果的条数，范围1-50，默认10
        # search_domain_filter="www.sohu.com",  # 只访问指定域名的内容
        search_recency_filter="noLimit",  # 搜索指定日期范围内的内容
        content_size="high"  # 控制网页摘要的字数，默认medium
    )
    print(response.search_result[0].content)
    return response.search_result[0].content


from datetime import datetime
from zhipuai import ZhipuAI


class Websearch:
    def __init__(self):
        # 初始化 ZhipuAI 客户端
        self.client = ZhipuAI(api_key=os.getenv("ZHIPU_KEY"))
        # 定义工具参数
        self.tools = [{
            "type": "web_search",
            "web_search": {
                "enable": True,
                "search_engine": "search_pro_sogou",  # 选择搜索引擎类型，默认为基础版
            }
        }]

    def __call__(self,user_input):
        date = datetime.now().strftime("%Y-%m-%d")

        # 系统提示模板，包含时间信息
        system_prompt = f"""你是一个具备网络访问能力的智能助手，在适当情况下，优先使用网络信息（参考信息）来回答，
        以确保用户得到最新、准确的帮助。当前日期是 {date}。"""

        # 用户输入的问题
        # user_input = "2025年美国总统是谁"

        # 构建动态用户问题提示
        user_question = f"参考最新消息给出对用户输入的回答: {user_input}"

        # 定义用户消息
        messages = [
            {
                "role": "system",
                "content": system_prompt
            },
            {
                "role": "user",
                "content": user_question
            }
        ]

        # 调用 API 获取响应
        response = self.client.chat.completions.create(
             # model="glm-4-air",  # 模型编码
            model="GLM-4.5-Flash",  # 模型编码
            messages=messages,  # 用户消息
            tools=self.tools  # 工具参数
        )

        print(response)
        # 打印响应结果
        print(response.choices[-1].message)
        return response.choices[-1].message.content

websearch = Websearch()


class RAG():
    def __init__(self):
        loader = TextLoader("/root/project/Code/sshcode/lc/character.txt",encoding="utf-8")
        docs = loader.load()

        model = getEmd()

        from langchain_text_splitters import RecursiveCharacterTextSplitter

        _splitter = RecursiveCharacterTextSplitter(chunk_size=100,chunk_overlap=20)

        split_doc = _splitter.split_documents(documents=docs)
        self.vector_db = Chroma.from_documents(split_doc,model)


    def __call__(self,query):
        _r = self.vector_db.similarity_search_with_score(query=query,k=2)
        return _r

from langchain_experimental.utilities import PythonREPL
repl = PythonREPL()


rag = RAG()



if __name__ == '__main__':
    # query = "2024年中国第一季度的GDP是多少"
    # result = get_web_search(query)
    # print(result)

    # web = Websearch()
    # r = web(query)
    # print(r)


    # query_rag = "林浅夏的父亲是干什么的"
    # r = rag(query_rag)
    # print(r)

    print(repl.run("print('123')"))
